#Anaconda/envs/pysagin python
# -*- coding: utf-8 -*-
'''
spacebase.py
============

.. module:: spacebase
  :platform: Windows, Linux
  :synopsis: 天基平台

.. moduleauthor:: WangXi

简介
----



版本
----

- 版本 1.0 (2025/03/27): 初始版本

'''

from astropy import units as u
from astropy.time import Time
from astropy.coordinates import CartesianRepresentation, ITRS, GCRS, EarthLocation, Angle
from poliastro.bodies import Earth
from poliastro.twobody import Orbit
from pysagin.device import Device, Radio, RadioFixedAccess, RadioRandomAccess
import numpy as np

from pysagin.base import Base

class SpaceBase(Base):
	def __init__(self, name:str, **kwargs):
		super().__init__(name,**kwargs)

class Satellite(SpaceBase):
    def __init__(self, name:str, a:float, e:float, inc:float, raan:float, argp:float, 
					init_nu:float, init_time:Time = Time("2025-03-27 00:00:00"), 
					**kwargs):
        """_summary_

		Args:
			name (str): 卫星名称
			a (float): 轨道半长轴,单位km
			e (float): 轨道偏心率,单位度
			inc (float): 轨道倾角,单位度
			raan (float): 轨道升交点经度,单位度
			argp (float): 轨道近地点幅角,单位度
			init_nu (float): 卫星初始真近点角,单位度
			init_time (astropy.time.Time): 卫星初始时间,默认Time("2025-03-27 00:00:00")
		"""
        super().__init__(name,**kwargs)
  
        self.name = name
        self.a = a
        self.e = e
        self.inc = inc
        self.raan = raan
        self.argp = argp
        self.init_nu = init_nu
        self.init_time = init_time
        self.current_time = init_time

        self.orb = Orbit.from_classical(attractor=Earth, a=self.a*u.km, ecc=self.e*u.one,
                                        inc=self._deg_to_rad(self.inc)*u.rad, 
                                        raan=self._deg_to_rad(self.raan)*u.rad, 
                                        argp=self._deg_to_rad(self.argp)*u.rad,
                                        nu=self._deg_to_rad(self.init_nu)*u.rad)
        
        self.update(self.init_time) # 触发首次更新以获得卫星的初始定位信息
        self.init_gps = self.current_gps
    
  
    def update(self, time) -> None:
        self.current_time = time
        propagated_orb = self.orb.propagate(self.current_time)

        # 获取卫星的位置向量
        r, _ = propagated_orb.rv()  # r: 位置向量, v: 速度向量

        # 将位置向量转换为地固坐标系 (ITRS)
        gcrs_position = GCRS(CartesianRepresentation(r),obstime=self.current_time)
        itrs_position = gcrs_position.transform_to(ITRS(obstime=self.current_time))

        # 将地固坐标系的位置转换为地理坐标
        location = EarthLocation(*itrs_position.cartesian.xyz)
        subpoint_lon, subpoint_lat, subpoint_alt = location.to_geodetic()

        self.current_gps = [subpoint_lon.value, subpoint_lat.value, subpoint_alt.value]
    
    def _deg_to_rad(self,angle_deg):
        """将角度转换为弧度制"""
        angle_rad = Angle(angle_deg, unit=u.deg).to(u.rad).to_value()
        # 确保角度在 -π 到 π 范围内
        return (angle_rad + np.pi) % (2 * np.pi) - np.pi

class Constellation:
    def __init__(self, name: str, type: str = "Walker", **kwargs):
        """
        初始化星座对象。

        Args:
            name (str): 星座名称。
            type (str): 星座类型，默认为 "Walker"
            **kwargs: 其他参数，用于定义 Walker 星座
                - T (int): 总卫星数量
                - P (int): 轨道面数量
                - F (float): 相位因子, 范围 [0, 1)
                - a (float): 半长轴, 单位km
                - e (float): 偏心率
                - inc (float): 轨道倾角, 单位度
                - init_time (astropy.time.Time): 初始时间，默认 Time("2025-03-27 00:00:00")
        """
        self.name = name
        self.type = type
        self.satellites: list[Satellite] = []  # 存储卫星对象列表
        self.links: list[list[str]] = []  # 存储星间链路信息

        if type.lower() == "walker":
            self.create_walker_constellation(**kwargs)
        else:
            raise ValueError(f"Unsupported constellation type: {type}")

    def create_walker_constellation(self, T: int, P: int, F: float, a: float, e: float,
                                    inc: float, init_time: Time = Time("2025-03-27 00:00:00")):
        """
        创建 Walker 星座

        Args:
            T (int): 总卫星数量
            P (int): 轨道面数量
            F (float): 相位因子, 范围 [0, 1)
            a (float): 半长轴, 单位km
            e (float): 偏心率
            inc (float): 倾角, 单位度
            init_time (astropy.time.Time): 初始时间
        """
        # 每个轨道面上的卫星数量
        S = T // P
        if T % P != 0:
            raise ValueError("Total satellites (T) must be divisible by the number of orbital planes (P).")

        for orbit_plane_num in range(P):
            raan = 360 / P * orbit_plane_num  # 升交点赤经（RAAN）
            for sat_num in range(S):
                # 计算真近点角（考虑相位因子）
                true_anomaly = 360 / S * sat_num + F * 360 / P * orbit_plane_num
                true_anomaly %= 360  # 确保角度在 [0, 360) 范围内

                # 创建卫星对象
                sat_name = f"{self.name}P{orbit_plane_num}S{sat_num}"
                satellite = Satellite(
                    name=sat_name,
                    a=a,
                    e=e,
                    inc=inc,
                    raan=raan,
                    argp=0,  # 近地点幅角设为 0
                    init_nu=true_anomaly,
                    init_time=init_time
                )
                self.satellites.append(satellite)

    def update_all(self, time: Time):
        """
        更新所有卫星的时间和卫星定位信息。

        Args:
            time (astropy.time.Time): 当前时间。
        """
        for satellite in self.satellites:
            satellite.update(time)

    def get_all_current_gps(self):
        """
        获取所有卫星的当前定位信息。

        Returns:
            list: 包含所有卫星的当前定位信息，每个位置是一个 [lon, lat, alt] 的列表。
        """
        return [satellite.get_current_gps() for satellite in self.satellites]
    
    def create_inter_satellite_links(self, Device_Radio:RadioFixedAccess, **kwargs):
        """
        在 Walker 星座中创建星间链路
        
        Args:
            Device_Radio (object): 所使用的无线电类名
            **kwargs: 无线电参数
        """
        if not hasattr(self, 'satellites') or len(self.satellites) == 0:
            raise ValueError("No satellites found in the constellation. Please create the constellation first.")

        T = len(self.satellites)  # 总卫星数量
        P = len({sat.raan for sat in self.satellites})  # 轨道面数量
        S = T // P  # 每个轨道面上的卫星数量

        # 将卫星按轨道面和相对位置组织成二维数组
        satellites_grid = [[None for _ in range(S)] for _ in range(P)]
        for sat in self.satellites:
            orbit_plane_num = int(sat.raan * P / 360)  # 根据 RAAN 确定轨道面编号
            sat_num = int(sat.init_nu * S / 360)  # 根据真近点角确定卫星编号
            satellites_grid[orbit_plane_num][sat_num] = sat

        links = []

        # 遍历每个轨道面和每颗卫星，创建星间链路
        for orbit_plane_num in range(P):
            for sat_num in range(S):
                current_sat:Satellite = satellites_grid[orbit_plane_num][sat_num]

                # 上邻居（同一轨道面上的前一颗卫星）
                upper_neighbor:Satellite = satellites_grid[orbit_plane_num][(sat_num - 1) % S]
                links.append([current_sat.name, upper_neighbor.name])

                # 下邻居（同一轨道面上的后一颗卫星）
                lower_neighbor:Satellite = satellites_grid[orbit_plane_num][(sat_num + 1) % S]
                links.append([current_sat.name, lower_neighbor.name])

                # 左邻居（相邻轨道面上相同相对位置的卫星）
                left_orbit_plane = (orbit_plane_num - 1) % P
                left_neighbor:Satellite = satellites_grid[left_orbit_plane][sat_num]
                links.append([current_sat.name, left_neighbor.name])

                # 右邻居（相邻轨道面上相同相对位置的卫星）
                right_orbit_plane = (orbit_plane_num + 1) % P
                right_neighbor:Satellite = satellites_grid[right_orbit_plane][sat_num]
                links.append([current_sat.name, right_neighbor.name])

        # 保存链路信息
        self.links = links
        # 创建星间链路设备
        self._create_inter_links_device(Device_Radio, **kwargs)
    
    def _create_inter_links_device(self, Device_Radio:RadioFixedAccess, **kwargs):
        for sat in self.satellites:
            device = Device_Radio(name=f"{sat.name}D{len(sat.devices)}", 
                                        platform=sat, **kwargs)
            sat.devices.append(device)
            
        for sat in self.satellites:
            for link in self.links:
                if sat.name == link[0]:
                    # 这里保证刚刚最新添加的无线电设备被相互连接
                    sat.devices[-1].add_access_radios([self.get_satallite_by_name(link[1]).devices[-1]])
    
    
    def get_satallite_by_name(self, name:str)->Satellite:
        """
        根据卫星名称获取卫星对象

        Args:
            name (str): 卫星名称

        Returns:
            Satellite: 卫星对象
        """
        for sat in self.satellites:
            if sat.name == name:
                return sat
        return None